#
# Copyright (C) 2023, Inria
# GRAPHDECO research group, https://team.inria.fr/graphdeco
# All rights reserved.
#
# This software is free for non-commercial, research and evaluation use 
# under the terms of the LICENSE.md file.
#
# For inquiries contact  george.drettakis@inria.fr
#
import torch
from einops import repeat
import torch.nn.functional as F
import math
from diff_gaussian_rasterization import GaussianRasterizationSettings, GaussianRasterizer
from scene.gaussian_model import GaussianModel
from scene.cameras import Camera


def generate_one_level_neural_gaussians(rasterizer: GaussianRasterizer, viewpoint_camera: Camera, pc: GaussianModel, is_training=False):
    visible_mask = rasterizer.markVisible(pc.get_anchor)
    anchor = pc.get_anchor[visible_mask]            # (N, 3)
    grid_scaling = pc.get_scaling[visible_mask]     # (N, 1)
    anchor_feat = pc.get_anchor_feat[visible_mask]
    
    # get view properties for anchor
    ob_view = anchor - viewpoint_camera.camera_center  # (N, 3)
    ob_dist = ob_view.norm(dim=1, keepdim=True)        # (N, 1)
    ob_view = ob_view / ob_dist                        # (N, 3)
    embed_view = pc.dir_embedder(ob_view) 
     
    cat_local_view = torch.cat([anchor_feat, embed_view], dim=1)
    neural_opacity = pc.get_opacity_mlp(cat_local_view) # [N, k]
    color = 2.0 * pc.get_color_mlp(cat_local_view).reshape(-1, 3)
    scale_rot = pc.get_cov_mlp(cat_local_view).reshape(-1, 7)
    offsets = pc.get_offset_mlp(anchor_feat).reshape(-1, 3)
    
    neural_opacity = neural_opacity.reshape(-1, 1)
    mask = (neural_opacity>0.0).view(-1)

    opacity = neural_opacity[mask]
            
    # combine for parallel masking
    concatenated = torch.cat([grid_scaling, anchor], dim=-1)
    concatenated_repeated = repeat(concatenated, 'n (c) -> (n k) (c)', k=pc.n_offsets)
    concatenated_all = torch.cat([concatenated_repeated, color, scale_rot, offsets], dim=-1)
    masked = concatenated_all[mask]
    
    repeat_scaling, repeat_anchor, color, scale_rot, offsets = masked.split([pc.scale_dim, 3, 3, 7, 3], dim=-1)
    if pc.scale_dim == 1 or pc.scale_dim == 3:
        repeat_offset_scaling = repeat_scaling
        repeat_scale_scaling = repeat_scaling
    elif pc.scale_dim == 2:
        repeat_offset_scaling = repeat_scaling[:, :1]
        repeat_scale_scaling = repeat_scaling[:, 1:]
    elif pc.scale_dim == 6:
        repeat_offset_scaling = repeat_scaling[:, :3]
        repeat_scale_scaling = repeat_scaling[:, 3:]
        
    xyz = repeat_anchor + offsets * repeat_offset_scaling
    
    rot = pc.rotation_activation(scale_rot[:, 3:7])
    scaling = torch.sigmoid(scale_rot[:, :3]) * repeat_scale_scaling
            
    if is_training:
        return xyz, color, opacity, scaling, rot, mask, visible_mask
    else:
        return xyz, color, opacity, scaling, rot


def render(viewpoint_camera, pc: GaussianModel, pipe, bg_color: torch.Tensor, scaling_modifier=1.0, retain_grad=False):
    """
        Render the scene. 
        Background tensor (bg_color) must be on GPU!
    """
    # Set up rasterization configuration
    tanfovx = math.tan(viewpoint_camera.FoVx * 0.5)
    tanfovy = math.tan(viewpoint_camera.FoVy * 0.5)

    raster_settings = GaussianRasterizationSettings(
        image_height=int(viewpoint_camera.image_height),
        image_width=int(viewpoint_camera.image_width),
        tanfovx=tanfovx,
        tanfovy=tanfovy,
        bg=bg_color,
        scale_modifier=scaling_modifier,
        viewmatrix=viewpoint_camera.world_view_transform,
        projmatrix=viewpoint_camera.full_proj_transform,
        sh_degree=1,
        campos=viewpoint_camera.camera_center,
        prefiltered=False,
        debug=pipe.debug
    )

    rasterizer = GaussianRasterizer(raster_settings=raster_settings)
    is_training = pc.get_color_mlp.training
    
    if is_training:
        xyz, color, opacity, scaling, rot, mask, visible_mask = generate_one_level_neural_gaussians(rasterizer, viewpoint_camera, pc, is_training)
    else:
        xyz, color, opacity, scaling, rot = generate_one_level_neural_gaussians(rasterizer, viewpoint_camera, pc, is_training)
        
    # Create zero tensor. We will use it to make pytorch return gradients of the 2D (screen-space) means
    screenspace_points = torch.zeros_like(xyz, dtype=xyz.dtype, requires_grad=True, device="cuda") + 0
    if retain_grad:
        try:    screenspace_points.retain_grad()
        except: pass
            
    # Rasterize visible Gaussians to image, obtain their radii (on screen). 
    rendered_image, radii = rasterizer(
        means3D = xyz,
        means2D = screenspace_points,
        shs = None,
        colors_precomp = color,
        opacities = opacity,
        scales = scaling,
        rotations = rot,
        cov3D_precomp = None)
    
    # Those Gaussians that were frustum culled or had a radius of 0 were not visible.
    if is_training:
        return {"render": rendered_image,
                "viewspace_points": screenspace_points,
                "visibility_filter" : radii > 0,
                "radii": radii,
                "selection_mask": mask,
                "neural_opacity": opacity,
                "scaling": scaling,
                "neural_xyz": xyz,
                "visible_mask": visible_mask
                }
    else:
        return {"render": rendered_image,
                "viewspace_points": screenspace_points,
                "visibility_filter" : radii > 0,
                "radii": radii}
